import imageio
import torch
from tqdm import tqdm
from animate import normalize_kp
from demo import load_checkpoints
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from skimage import img_as_ubyte
from skimage.transform import resize
import cv2
import os
import argparse
import subprocess
import shutil
import sys
from tkinter import *
import cv2
from PIL import Image, ImageTk#图像控件
import threading#多线程
from tkinter import filedialog
from face import puts_
from funtion import vector_centerpoint
from funtion import similarity
from funtion import video_brodcast
from image_animation import main, func1
from openpose import Init2, func2
import math
from os import startfile, system
from moviepy.editor import*
import mp3play
# -----------------------窗口----------------------------

window = Tk()
window.title("这是一个窗口")
window.geometry("1200x550")
# window.overrideredirect(1)#删除标题栏
window.resizable(0,0) #禁止拉伸

#定义画布
canvas=Canvas(window,width=600,height=600)
canvas.place(x=0,y=0)

canvas1=Canvas(window,width=600,height=600)
canvas1.place(x=600,y=0)

#--------------------------------------------------------

# ----------------------输入-----------------------------

# 构造输入信息结构
ap = argparse.ArgumentParser()
# 静态图片地址
ap.add_argument("-i", "--input_image", required=True, help="Path to image to animate")
ap.add_argument("-j", "--input_image1", required=True, help="Path to image to animate")
ap.add_argument("-m", "--input_image2", required=True, help="Path to image to animate")
# 权重、模型地址
ap.add_argument("-c", "--checkpoint", required=True, help="Path to checkpoint")
# 被捕捉的视频地址
ap.add_argument("-v", "--input_video", required=False, help="Path to video input")

#--------------------------------------------------------

# ---------------------读取mp3---------------------------



#--------------------------------------------------------

# ----------------按键逻辑初始化--------------------------

change_click=0
img_change_click=0
remake_click=0
exit_click=0
# #切换
def l():
    global change_click
    change_click+=1
#换头像
def k():
    global img_change_click
    img_change_click+=1
#校准
def r():
    global remake_click
    remake_click+=1
#退出
def q():
    global exit_click
    exit_click+=1

#--------------------------------------------------------

# -------------------image_animation初始化----------------

def Init1(ap, idx):
    V_list, cv2_source1, cap = main(ap, idx)
    count = 0
    tmp = []
    return [V_list, cv2_source1, cap, count, tmp]

#--------------------------------------------------------
movie_play=1
movie_path='doudou2.mp4'
capture1 = cv2.VideoCapture('doudou2.mp4')

#-----------------------功能模块--------------------------
def video_demo1():
    def cc():
        idx = 0
        function = 0
        function2=0
        V_list1, cv2_source1, cap, count, tmp = Init1(ap, idx)
        V_list2 = Init2()
        global change_click,img_change_click,remake_click,exit_click
        #openpose识别姿态模块：
        while True:
            if function == 0:
                while True:
                    frame,flag= func2(cap, V_list2)
                    frame=cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    current_img2 = Image.fromarray(np.uint8(frame)).resize((600,500))
                    imagetk2 = ImageTk.PhotoImage(current_img2)
                    canvas.create_image(0,0,anchor='nw',image=imagetk2)
                    canvas.image = imagetk2
                    if(flag):
                        global movie_play,movie_path,capture1
                        movie_play=1
                        movie_path = filedialog.askopenfilename()#打开文件路径
                        capture1 = cv2.VideoCapture(movie_path)#获取视频
                        # video = VideoFileClip(movie_path)
                        # audio = video.audio
                        function2=1
                        break
                        # audio = capture1.audio
                        # audio.write_audiofile('music.mp3')

                    #切换openpose&image_animation模块
                    if change_click:
                            function ^= 1
                            V_list1, cv2_source1, cap, count, tmp = Init1(ap, idx)
                            change_click=0
                    #关闭所有窗口
                    if exit_click:
                        exit_click=0
                        window.destroy()
                    #判断openpose姿态识别退出
                    if(function==1):
                        break
            print(function,end='  ')
            print(function2)
            if (function == 0) and  (function2==1):
                print(1)
            #image_animation：人脸控制模块
            if function:
                while True:
                    cv2_source, im, frame1, out1, tmp = func1(V_list1, count, tmp, cv2_source1)
                    count += 1
                    frame1 = cv2.cvtColor(frame1.astype('float32'), cv2.COLOR_BGR2RGB)#转换颜色从BGR到RGB
                    im=cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
                    current_img = Image.fromarray(np.uint8(frame1*255)).resize((600,500))
                    current_img1 = Image.fromarray(np.uint8(im*255)).resize((600,500))#图像从array转换成Image对象，并修改尺寸
                    imagetk = ImageTk.PhotoImage(current_img)
                    imagetk1 = ImageTk.PhotoImage(current_img1)#将图像转换成与tkinter兼容的照片
                    canvas.create_image(0,0,anchor='nw',image=imagetk)
                    canvas1.create_image(0,0,anchor='nw',image=imagetk1)#将图像更新到画布上
                    canvas.image = imagetk
                    canvas1.image = imagetk1#防闪烁

                    #切换openpose&image_animation模块
                    if change_click:
                            function ^= 1
                            V_list1, cv2_source1, cap, count, tmp = Init1(ap, idx)
                            change_click=0
                            break
                    #换脸
                    if img_change_click:
                        idx ^= 1
                        V_list1, cv2_source1, cap, count, tmp = Init1(ap, idx)
                        img_change_click=0
                    #重新读取关键点
                    if remake_click:
                        V_list1, cv2_source1, cap, count, tmp = Init1(ap, idx)
                        remake_click=0
                    #关闭所有窗口
                    if exit_click:
                        exit_click=0
                        window.destroy()
                    #判断image_animation退出
                    if(function==0):
                        break
            #openpose播放视频模块
            if (function == 0) and (function2==1):
                f= mp3play.load('music.mp3')
                f.play()
                function2=0
                while True:
                    for i in range(3):
                        ret,frame1 = capture1.read()#从视频文件读取照片
                    if ret:
                        frame1 = cv2.cvtColor(frame1, cv2.COLOR_BGR2RGB)#转换颜色从BGR到RGB
                        current_img1 = Image.fromarray(frame1).resize((600,500))#图像从array转换成Image对象，并修改尺寸
                        imagetk1 = ImageTk.PhotoImage(current_img1)#将图像转换成与tkinter兼容的照片
                        canvas1.create_image(0,0,anchor='nw',image=imagetk1)#将图像更新到画布上
                        canvas1.image = imagetk1
                    
                    #关闭所有窗口
                    if exit_click:
                        exit_click=0
                        window.destroy()
                    #判断openpose退出
                    if(function==1):
                        break
    threading.Thread(target=cc).start()

video_demo1()

#--------------------------------------------------------

# ----------------按钮界面设置----------------------------

bt_l= Button(window,text='切换',height=2,width=40,command=l)
bt_l.place(x=0,y=500)

bt_k= Button(window,text='换头',height=2,width=40,command=k)
bt_k.place(x=300,y=500)

bt_r= Button(window,text='校准',height=2,width=40,command=r)
bt_r.place(x=600,y=500)

bt_q = Button(window,text='退出',height=2,width=40,command=q)
bt_q.place(x=900,y=500)

window.mainloop()

#--------------------------------------------------------